from confluent_kafka import Consumer, Producer
import logging
import json
from cryptography.fernet import Fernet
from datetime import datetime
import hmac
import hashlib
from .secure_config import SecureConfigLoader

class LegacyPayloadAdapter:
    """处理遗留二进制格式的适配器"""
    HEADER_FORMATS = {
        b'\x00\x01': 'avro-binary-v1',
        b'\x00\x02': 'protobuf-legacy',
        b'\x00\x03': 'thrift-encoded'
    }

    def __init__(self, schema_registry_url):
        self.schema_registry_url = schema_registry_url

    def decode(self, payload):
        """解码遗留二进制格式"""
        header = payload[:2]
        fmt = self.HEADER_FORMATS.get(header, 'unknown')
        
        if fmt == 'avro-binary-v1':
            return self._decode_avro(payload[2:])
        elif fmt == 'protobuf-legacy':
            return self._decode_protobuf(payload[2:])
        else:
            raise ValueError(f"Unsupported format: {fmt}")

    def _decode_avro(self, payload):
        # 实现Avro解码逻辑
        return {'format': 'avro', 'data': payload.hex()}

    def _decode_protobuf(self, payload):
        # 实现Protobuf解码逻辑  
        return {'format': 'protobuf', 'data': payload.hex()}

class KafkaAuditor:
    def __init__(self, config_loader):
        self.config_loader = config_loader
        self.logger = logging.getLogger(__name__)
        self.legacy_adapter = LegacyPayloadAdapter(
            config_loader.get('schema.registry.url')
        )
        
        # 初始化加密组件
        self.cipher = Fernet(config_loader.secret_key)
        
        # 创建审计日志生产者
        producer_conf = {
            'bootstrap.servers': config_loader.get('bootstrap.servers'),
            'security.protocol': config_loader.get('security.protocol'),
            'ssl.ciphers': config_loader.get('ssl.ciphers'),
            'ssl.min.version': config_loader.get('ssl.min.version')
        }
        self.audit_producer = Producer(producer_conf)

    def _generate_audit_record(self, event_type, metadata):
        """生成加密审计记录"""
        record = {
            'timestamp': datetime.utcnow().isoformat(),
            'event_type': event_type,
            'metadata': metadata,
            'signature': hmac.new(
                self.config_loader.secret_key,
                msg=str(metadata).encode(),
                digestmod=hashlib.sha512
            ).hexdigest()
        }
        return self.cipher.encrypt(json.dumps(record).encode())

    def audit_produce(self, topic, value, key=None):
        """审计消息生产"""
        try:
            # 处理遗留格式
            if isinstance(value, bytes) and value[:2] in LegacyPayloadAdapter.HEADER_FORMATS:
                decoded = self.legacy_adapter.decode(value)
                metadata = {
                    'topic': topic,
                    'format': 'legacy',
                    'decoded': decoded
                }
            else:
                metadata = {
                    'topic': topic,
                    'format': 'modern',
                    'size': len(value)
                }

            # 发送审计记录
            audit_record = self._generate_audit_record('produce', metadata)
            self.audit_producer.produce(
                'kafka-audit-log',
                value=audit_record
            )
            
        except Exception as e:
            self.logger.error(f"审计记录失败: {str(e)}")

    def audit_consume(self, msg):
        """审计消息消费"""
        try:
            metadata = {
                'topic': msg.topic(),
                'partition': msg.partition(),
                'offset': msg.offset(),
                'key': msg.key().decode() if msg.key() else None
            }
            
            audit_record = self._generate_audit_record('consume', metadata)
            self.audit_producer.produce(
                'kafka-audit-log',
                value=audit_record
            )
            
        except Exception as e:
            self.logger.error(f"审计记录失败: {str(e)}")

    def flush_audits(self):
        """刷新审计记录"""
        self.audit_producer.flush()

if __name__ == '__main__':
    # 示例用法
    config_loader = SecureConfigLoader(
        config_path='/etc/app/kafka.conf.enc',
        secret_key=os.getenv('CONFIG_SECRET')
    )
    auditor = KafkaAuditor(config_loader)
    
    # 模拟生产审计
    auditor.audit_produce('air-quality', b'\x00\x01testdata')
    auditor.flush_audits()